/* LTO plugin for gold and/or GNU ld.
   Copyright (C) 2009-2019 Free Software Foundation, Inc.
   Contributed by Rafael Avila de Espindola (espindola@google.com).

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 3, or (at your option)
any later version.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; see the file COPYING3.  If not see
<http://www.gnu.org/licenses/>.  */

/* The plugin has only one external function: onload. Gold passes it an array of
   function that the plugin uses to communicate back to gold.

   With the functions provided by gold, the plugin can be notified when
   gold first analyzes a file and pass a symbol table back to gold. The plugin
   is also notified when all symbols have been read and it is time to generate
   machine code for the necessary symbols.

   More information at http://gcc.gnu.org/wiki/whopr/driver.

   This plugin should be passed the lto-wrapper options and will forward them.
   It also has options at his own:
   -debug: Print the command line used to run lto-wrapper.
   -nop: Instead of running lto-wrapper, pass the original to the plugin. This
   only works if the input files are hybrid. 
   -linker-output-known: Do not determine linker output
   -sym-style={none,win32,underscore|uscore}
   -pass-through  */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif
#if HAVE_STDINT_H
#include <stdint.h>
#endif
#include <assert.h>
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <inttypes.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#ifdef HAVE_SYS_WAIT_H
#include <sys/wait.h>
#endif
#ifndef WIFEXITED
#define WIFEXITED(S) (((S)&0xff) == 0)
#endif
#ifndef WEXITSTATUS
#define WEXITSTATUS(S) (((S)&0xff00) >> 8)
#endif
#include <filenames.h>
#include <hashtab.h>
#include "simple-object.h"
#include "plugin-api.h"

#include <obstack.h>
#define obstack_chunk_alloc malloc
#define obstack_chunk_free free

/* We need to use I64 instead of ll width-specifier on native Windows.
   The reason for this is that older MS-runtimes don't support the ll.  */
#ifdef __MINGW32__
#define PRI_LL "I64"
#else
#define PRI_LL "ll"
#endif

/* Handle opening elf files on hosts, such as Windows, that may use
   text file handling that will break binary access.  */
#ifndef O_BINARY
#define O_BINARY 0
#endif
namespace LIBIBERTY
{
#include <libiberty.h>
}
using LIBIBERTY::xmalloc;
using LIBIBERTY::lbasename;
using LIBIBERTY::xstrdup;
using LIBIBERTY::getpwd;
using LIBIBERTY::concat;
using LIBIBERTY::lrealpath;

#include <vector>
#include <string>
#include <map>

using std::map;
using std::string;
using std::vector;

#if HAVE_DOS_BASSED_FILE_SYSTEM
const char *separator = "\\";
#else
const char *separator = "/";
#endif

#define DEFAULT_BOLT_OUT_DIR (getpwd())
#define DEFAULT_BOLT_OUT_NAME "default.fbolt"
#define DEFAULT_BOLT_OUT_NAME_SUFFIX ".fbolt"

/* bolt magic section name.  */

#define BOLT_SECTION_PREFIX ".gnu.bolt.text"
#define BOLT_SECTION_CALLER_PREFIX ".bolt.caller"
#define BOLT_SECTION_CALLEE_PREFIX ".bolt.callee"

extern "C"{
  enum ld_plugin_status onload(struct ld_plugin_tv *tv);
}
/* Encapsulates object file data during symbol scan.  */
struct plugin_objfile
{
  int found;
  int offload;
  simple_object_read *objfile;
  const struct ld_plugin_input_file *file;
};

// struct gcc_options_info
// {
//   int argc;
//   char * argv;
// };

// struct gcc_options_info gcc_options

vector<string> gcc_options;

struct jump_info
{
  string des_func_name;
  string src_addr_offset;
  string des_addr_offset;
  string count;
  int valid = 0;
};

struct edge_info
{
  int index;
  int file_num = -1;
  vector<struct jump_info> jumps;
};

map<string, std::vector<struct edge_info>>
    all_functions;
vector<struct ld_plugin_section> sections_list;
vector<struct ld_plugin_section> bolt_sections_list;
static int g_file_num = -1;

static char *arguments_file_name = NULL;
static ld_plugin_register_claim_file register_claim_file = NULL;
static ld_plugin_register_all_symbols_read register_all_symbols_read = NULL;
static ld_plugin_register_cleanup register_cleanup = NULL;
static ld_plugin_message message = NULL;

static ld_plugin_update_section_order update_section_order = NULL;
static ld_plugin_allow_section_ordering allow_section_ordering = NULL;
static ld_plugin_get_input_section_count get_input_section_count = NULL;
static ld_plugin_get_input_section_type get_input_section_type = NULL;
static ld_plugin_get_input_section_name get_input_section_name = NULL;

static enum ld_plugin_output_file_type linker_output;

static int linker_output_set;
static int linker_output_known;

/* The version of gold being used, or -1 if not gold.  The number is
   MAJOR * 100 + MINOR.  */
static int gold_version = -1;

static const int BUFFLEN = 1024;

static string out_file_name = NULL;
static string bolt_dir_path = NULL;
static FILE *bolt_file_fd = NULL;

static void
check_1(int gate, enum ld_plugin_level level, const char *text)
{
  if (gate)
    return;

  if (message)
    message(level, text);
  else
  {
    /* If there is no nicer way to inform the user, fallback to stderr. */
    fprintf(stderr, "%s\n", text);
    if (level == LDPL_FATAL)
      abort();
  }
}

/* This little wrapper allows check to be called with a non-integer
   first argument, such as a pointer that must be non-NULL.  We can't
   use c99 bool type to coerce it into range, so we explicitly test.  */
#define check(GATE, LEVEL, TEXT) check_1(((GATE) != 0), (LEVEL), (TEXT))

#define __MSG_INFO__
#define __MSG_WARN__
#define __MSG_ERROR__

#ifdef __MSG_INFO__
#define MSG_INFO(...)                                            \
  if (message)                                                   \
  {                                                              \
    message(LDPL_INFO, "-- [bolt-plugin] INFO -- " __VA_ARGS__); \
  }                                                              \
  else                                                           \
  {                                                              \
    fprintf(stderr, "---- [bolt-plugin] INFO -- " __VA_ARGS__);  \
  }
#else
#define MSG_INFO(...)
#endif

#ifdef __MSG_WARN__
#define MSG_WARN(...)                                               \
  if (message)                                                      \
  {                                                                 \
    message(LDPL_WARNING, "-- [bolt-plugin] WARN -- " __VA_ARGS__); \
  }                                                                 \
  else                                                              \
  {                                                                 \
    fprintf(stderr, "-- [bolt-plugin] WARN -- " __VA_ARGS__);       \
  }
#else
#define MSG_WARN(...)
#endif

#ifdef __MSG_ERROR__
#define MSG_ERROR(...)                                             \
  if (message)                                                     \
  {                                                                \
    message(LDPL_FATAL, "-- [bolt-plugin] ERROR -- " __VA_ARGS__); \
  }                                                                \
  else                                                             \
  {                                                                \
    fprintf(stderr, "-- [bolt-plugin] ERROR -- " __VA_ARGS__);     \
    abort();                                                       \
  }
#else
#define MSG_ERROR(...)
#endif

inline static int is_prefix_of(const char *prefix, const char *str)
{
  return strncmp(prefix, str, strlen(prefix)) == 0;
}

/* Remove temporary files at the end of the link. */

static enum ld_plugin_status
cleanup_handler(void)
{
  if (bolt_file_fd)
  {
    fclose(bolt_file_fd);
  }
  return LDPS_OK;
}

static void
open_file(const char *file_name)
{
  if (file_name == NULL)
  {
    MSG_ERROR("Unable to build .fbolt file, maybe indeed \"-fauto-bolt=*\"");
  }

  if (bolt_file_fd == NULL)
  {
    MSG_INFO("Generate file for bolt: %s", file_name);
    bolt_file_fd = fopen(file_name, "wt");
    if (!bolt_file_fd)
    {
      MSG_ERROR("Please check whether the target path exists. Failed to open the file: %s.", file_name)
    }
    return;
  }
  else
  {
    MSG_WARN("File %s is open. Skip here", file_name);
  }
}

static void dump_to_bolt_file()
{
  string fnname;
  int caller_len;
  int callee_len;
  int i;
  int j;

  if (!bolt_file_fd)
  {
    open_file(out_file_name.c_str());
  }

  for (auto iter : all_functions)
  {
    caller_len = iter.second.size();

    for (i = 0; i < caller_len; i++)
    {
      fnname = (caller_len == 1) ? iter.first : iter.first + "/" + std::to_string(i + 1);
      for (auto jump : iter.second[i].jumps)
      {
        if (jump.des_func_name.empty())
        {
          jump.des_func_name = fnname;
        }
        else
        {
          callee_len = all_functions[jump.des_func_name].size();

          if (callee_len > 1)
          {
            bool found = false;
            for (j = 0; j < callee_len; j++)
            {
              if (all_functions[jump.des_func_name][j].file_num == iter.second[i].file_num)
              {
                check(all_functions[jump.des_func_name][j].file_num != -1, LDPL_FATAL,
                      "file_num is invalid!");
                jump.des_func_name = jump.des_func_name + "/" + std::to_string(j + 1);
                found = true;
                break;
              }
            }

            if (!found)
            {
              continue;
            }
          }
        }
        fprintf(bolt_file_fd, "1 %s %s 1 %s %s 0 %s\n", fnname.c_str(), jump.src_addr_offset.c_str(),
                jump.des_func_name.c_str(), jump.des_addr_offset.c_str(), jump.count.c_str());
      }
    }
    fflush(bolt_file_fd);
  }
}

static void reorder_funcs()
{
  int len = sections_list.size();

  if (len > 0)
  {
    MSG_INFO("reorder_funcs");

    struct ld_plugin_section *section_list = NULL;
    section_list = (struct ld_plugin_section *)xmalloc(len*sizeof(struct ld_plugin_section));

    memcpy(section_list, &sections_list[0], len * sizeof(struct ld_plugin_section));
    update_section_order(section_list, len);

    free(section_list);
  }
}

/* Called by the linker once all symbols have been read. */

static enum ld_plugin_status
all_symbols_read_handler(void)
{
  if (!all_functions.empty())
  {
    if (gold_version != -1)
    {
      if (!bolt_sections_list.empty())
      {
        for (auto section : bolt_sections_list)
        {
          MSG_INFO("inset_null_function")
          {
            sections_list.push_back(section);
          }
        }
      }
      reorder_funcs();
    }
    dump_to_bolt_file;
  }

  return LDPS_OK;
}

static char *
get_next_content(char *p, char *end)
{
  while (*p && p < end)
  {
    p++;
  }
  p++;
  return p;
}
/* Translate the IL symbol table located between DATA and END. Append the
   slots and symbols to OUT. */

static void
translate(char *data, char *end)
{
  struct edge_info bb_edge;
  string fnname;
  string valid;

  if (!is_prefix_of(BOLT_SECTION_CALLER_PREFIX, data))
  {
    return;
  }

  check(data, LDPL_FATAL, "Caller is NULL");
  fnname = strdup(data + strlen(BOLT_SECTION_CALLER_PREFIX));

  data = get_next_content(data, end);
  check(data, LDPL_FATAL, "data is NULL");

  while (*data && data < end)
  {
    struct jump_info jump;

    jump.src_addr_offset = strdup(data);
    check(data, LDPL_FATAL, "data is NULL");

    if (is_prefix_of(BOLT_SECTION_CALLEE_PREFIX, data))
    {
      jump.des_func_name = strdup(data + strlen(BOLT_SECTION_CALLEE_PREFIX));
      jump.des_addr_offset = "0";
    }
    else
    {
      jump.des_addr_offset = strdup(data);
    }

    data = get_next_content(data, end);
    check(data, LDPL_FATAL, "data is NULL");

    jump.count = strdup(data);

    bb_edge.jumps.push_back(jump);

    data = get_next_content(data, end);
    check(data, LDPL_FATAL, "data is NULL");
  }
  if (!fnname.empty())
  {
    auto iter = all_functions.find(fnname);
    if (iter != all_functions.end())
    {
      bb_edge.index = iter->second.size() + 1;
    }
    else
    {
      bb_edge.index = 1;
    }
    bb_edge.file_num = g_file_num;

    all_functions[fnname].push_back(bb_edge);
  }
}

/* Process one section of an object file.  */

static int
process_symtab(void *data, const char *name, off_t offset, off_t length)
{
  struct plugin_objfile *obj = (struct plugin_objfile *)data;
  char *s;
  char *secdatastart;
  char *secdata;

  if (!is_prefix_of(BOLT_SECTION_PREFIX, name))
  {
    return 1;
  }

  secdata = secdatastart = (char *)xmalloc(length*sizeof(char));
  offset += obj->file->offset;
  if (offset != lseek(obj->file->fd, offset, SEEK_SET))
    goto err;

  do
  {
    ssize_t got = read(obj->file->fd, secdata, length);
    if (got == 0)
      break;
    else if (got > 0)
    {
      secdata += got;
      length -= got;
    }
    else if (errno != EINTR)
      goto err;
  } while (length > 0);
  if (length > 0)
    goto err;

  translate(secdatastart, secdata);
  obj->found++;
  free(secdatastart);
  return 1;

err:
  if (message)
    message(LDPL_FATAL, "%s: corrupt object file", obj->file->name);
  /* Force claim_file_handler to abandon this file.  */
  obj->found = 0;
  free(secdatastart);
  return 0;
}

string section_types[] = {".text.startup.", ".text.hot.", ".text.unlikely.", ".text."}; 
void record_function_sections(const struct ld_plugin_input_file *file)
{
  unsigned int count = 0;
  struct ld_plugin_section section;
  unsigned int shndx;
  char *name = NULL;

  (*allow_section_ordering)();

  (*get_input_section_count)(file->handle, &count);

  for (shndx = 0; shndx < count; ++shndx)
  {
    section.handle = file->handle;
    section.shndx = shndx;

    if (get_input_section_name(section, &name))
    {
      break;
    }

    for (auto i : section_types)
    {
      if (is_prefix_of(i.c_str(), name))
      {
        if (!is_prefix_of("test_bolt.o", lbasename(file->name)))
        {
          sections_list.push_back(section);
        }
        else
        {
          bolt_sections_list.push_back(section);
        }
        break;
      }
    }
  }
}

/* Callback used by gold to check if the plugin will claim FILE. Writes
   the result in CLAIMED. */

static enum ld_plugin_status
claim_file_handler(const struct ld_plugin_input_file *file, int *claimed)
{
  enum ld_plugin_status status;
  struct plugin_objfile obj;
  int err;
  const char *errmsg;

  *claimed = 0;

  if (gold_version != -1)
  {
    record_function_sections(file);
  }

  obj.file = file;
  obj.found = 0;
  obj.offload = 0;

  errmsg = NULL;
  obj.objfile = simple_object_start_read(file->fd, file->offset, NULL,
                                         &errmsg, &err);

  /* No file, but also no error code means unrecognized format; just skip it.  */
  if (!obj.objfile && !err){
    MSG_ERROR("%s: %s", file->name, err);
  }
    
  if (obj.objfile){
    errmsg = simple_object_find_sections(obj.objfile, process_symtab, &obj, &err);
  }

  if (!obj.objfile || errmsg)
  {
    MSG_ERROR("%s: %s", file->name, errmsg);
  }

  goto cleanup;

cleanup:
  if (obj.objfile)
    simple_object_release_read(obj.objfile);

  return LDPS_OK;
}

/* Mangle filename path of BASE and output new allocated pointer with
   mangled path.  */

char *
mangle_path(char const *base)
{
  /* Convert '/' to '#', convert '..' to '^',
     convert ':' to '~' on DOS based file system.  */
  const char *probe;
  char *buffer = (char *)xmalloc(strlen(base) + 1);
  char *ptr = buffer;

#if HAVE_DOS_BASED_FILE_SYSTEM
  if (base[0] && base[1] == ':')
  {
    ptr[0] = base[0];
    ptr[1] = '~';
    ptr += 2;
    base += 2;
  }
#endif
  for (; *base; base = probe)
  {
    size_t len;

    for (probe = base; *probe; probe++)
      if (*probe == '/')
        break;
    len = probe - base;
    if (len == 2 && base[0] == '.' && base[1] == '.')
      *ptr++ = '^';
    else
    {
      memcpy(ptr, base, len);
      ptr += len;
    }
    if (*probe)
    {
      *ptr++ = '#';
      probe++;
    }
  }

  /* Terminate the string.  */
  *ptr = '\0';

  return buffer;
}

static string
prefix_out_file_name(string file_name)
{
  if (!IS_ABSOLUTE_PATH(file_name.c_str()))
  {
    if (!bolt_dir_path.empty())
    {
      file_name = concat(getpwd(), separator, file_name.c_str(), NULL) ;
      file_name = mangle_path(file_name.c_str());
    }
    else
    {
      bolt_dir_path = DEFAULT_BOLT_OUT_DIR;
    }
  }
  file_name = concat(bolt_dir_path.c_str(), separator, file_name.c_str(), DEFAULT_BOLT_OUT_NAME_SUFFIX, NULL);
  return file_name;
}

static int
match_gcc_option(const char *option_prefix)
{
  if (option_prefix == NULL)
  {
    return -1;
  }
  for (int i = 0; i < gcc_options.size(); i++)
  {
    if (is_prefix_of(option_prefix, gcc_options[i].c_str()))
    {
      return i;
    }
  }
  return -1;
}

static void
get_options_form_collect_gcc_options(const char *collect_gcc, const char *collect_gcc_options)
{
  if (collect_gcc == NULL || collect_gcc_options == NULL)
  {
    return;
  }

  int len = strlen(collect_gcc_options);
  unsigned int l = 0;
  unsigned int r = 0;
  while(r < len && collect_gcc_options[r] != '\0')
  {
    if (collect_gcc_options[r] == '\'')
    {
      string option;
      ++r;
      do
      {
        if (collect_gcc_options[r] == '\0')
        {
          MSG_ERROR("malformed COLLECT_GCC_OPTIONS");
        }
        else if (is_prefix_of("'\\''", &collect_gcc_options[r]))
        {
          option.push_back('\'');
          r += 4;
        }
        else if (collect_gcc_options[r] == '\'')
        {
          break;
        }
        else
        {
          option.push_back(collect_gcc_options[r]);
          r++;
        }
      } while (1);
      MSG_INFO("option: %s", option.c_str());
      gcc_options.push_back(option);
    }
  }
}

/* Parse the plugin options. */

static void
process_gcc_option()
{
  string flag_auto_bolt_ = "-fauto-bolt=";
  string flag_o = "-o";

  char *collect_gcc = getenv("COLLECT_GCC");
  char *collect_gcc_option = getenv("COLLECT_GCC_OPTION");

  get_options_form_collect_gcc_options(collect_gcc, collect_gcc_option);

  const int auto_bolt_index = match_gcc_option(flag_auto_bolt_.c_str());
  if (auto_bolt_index != -1)
  {
    bolt_dir_path = lrealpath(gcc_options[auto_bolt_index].substr(flag_auto_bolt_.size()).c_str());
  }
  else
  {
    MSG_WARN("Output directory is not set, you may need to use \"-fauto-bolt=*\", default is %s.", DEFAULT_BOLT_OUT_DIR);
  }

  const int o_index = match_gcc_option(flag_o.c_str());
  if (o_index != -1)
  {
    out_file_name = prefix_out_file_name(gcc_options[o_index + 1]);
  }
  else
  {
    out_file_name = DEFAULT_BOLT_OUT_NAME;
    MSG_WARN("The output file name is not set, you may need to use \"-o\", default is %s.", DEFAULT_BOLT_OUT_NAME);
  }
}

static void
init_plugin()
{
  process_gcc_option();
}

static void
process_option(const char * options){

}

/* Called by gold after loading the plugin. TV is the transfer vector. */

enum ld_plugin_status
onload(struct ld_plugin_tv *tv)
{
  struct ld_plugin_tv *p;
  enum ld_plugin_status status;

  p = tv;
  while (p->tv_tag)
  {
    switch (p->tv_tag)
    {
    case LDPT_MESSAGE:
      message = p->tv_u.tv_message;
      break;
    case LDPT_REGISTER_CLAIM_FILE_HOOK:
      register_claim_file = p->tv_u.tv_register_claim_file;
      break;
    case LDPT_REGISTER_ALL_SYMBOLS_READ_HOOK:
      register_all_symbols_read = p->tv_u.tv_register_all_symbols_read;
      break;
    case LDPT_REGISTER_CLEANUP_HOOK:
      register_cleanup = p->tv_u.tv_register_cleanup;
      break;
    case LDPT_OPTION:
      process_option(p->tv_u.tv_string);
      break;
    case LDPT_GOLD_VERSION:
      gold_version = p->tv_u.tv_val;
      break;
    case LDPT_LINKER_OUTPUT:
      linker_output = (enum ld_plugin_output_file_type)p->tv_u.tv_val;
      linker_output_set = 1;
      break;
    case LDPT_ALLOW_SECTION_ORDERING:
      allow_section_ordering = p->tv_u.tv_allow_section_ordering;
      break;
    case LDPT_UPDATE_SECTION_ORDER:
      update_section_order = p->tv_u.tv_update_section_order;
      break;
    case LDPT_GET_INPUT_SECTION_NAME:
      get_input_section_name = p->tv_u.tv_get_input_section_name;
      break;
    case LDPT_GET_INPUT_SECTION_COUNT:
      get_input_section_count = p->tv_u.tv_get_input_section_count;
      break;
    default:
      break;
    }
    p++;
  }

  if (linker_output_set && linker_output != LDPO_EXEC)
  {
    MSG_INFO("This linker[%d] is not for exec, just skip.", linker_output);
  }

  check(register_claim_file, LDPL_FATAL, "register_claim_file not found");
  status = register_claim_file(claim_file_handler);
  check(status == LDPS_OK, LDPL_FATAL,
        "could not register the claim_file callback");

  if (register_cleanup)
  {
    status = register_cleanup(cleanup_handler);
    check(status == LDPS_OK, LDPL_FATAL,
          "could not register the cleanup callback");
  }

  if (register_all_symbols_read)
  {
    status = register_all_symbols_read(all_symbols_read_handler);
    check(status == LDPS_OK, LDPL_FATAL,
          "could not register the all_symbols_read callback");
  }

  if (gold_version != -1)
  {
    MSG_INFO("Using gold linker.");
    check(allow_section_ordering, LDPL_FATAL, "allow_section_ordering not found");
    check(update_section_order, LDPL_FATAL, "update_section_order not found");
    check(get_input_section_name, LDPL_FATAL, "get_input_section_name not found");
    check(get_input_section_count, LDPL_FATAL, "get_input_section_count not found");
  }

  init_plugin();

  return LDPS_OK;
}
